import java.util.Arrays;

public class Solution {
    public static void main(String[] args) {

    }

    public int minMalwareSpread(int[][] graph, int[] initial) {
        /**
         * 尽量减少恶意软件的传播Ⅱ
         * 很复杂，以我现在知识学这个还是太勉强了*/
        // 1 预处理
        int n = graph.length;
        initialize(n, initial);

        // 2 合并所有非感染节点
        for(int i = 0; i < n; i++) {
            if(infect[i]) {
                continue;
            }
            for(int j = i+1; j < n; j++) {
                if(infect[j]) {
                    continue;
                }
                if(graph[i][j] == 1) {
                    union(i, j);
                }
            }
        }

        // 3 打标签（遍历所有感染节点）
        for(int i: initial) {
            for(int j = 0; j < n; j++) {
                // -两点不能相同 且 另一点不能为病毒，即保证都是普通点 且 要保证两点有连接
                if(i!=j && !infect[j] && graph[i][j]==1) {
                    int fj = find(j);
                    if(representative[fj] == -1 || representative[fj] == i) {
                        // -集合没有感染，或感染元素为同一个
                        representative[fj] = i;
                    } else {
                        representative[fj] = -2;
                    }
                }
            }
        }

        // 4 统计（一个病毒可能还会感染多个集合）
        int[] count = new int[n];
        for(int i = 0; i < n; i++) {
            if(representative[find(i)] >= 0 && find(i) == i) {
                count[representative[find(i)]] += size[find(i)];
            }
        }

        // 5 返回值
        int minIndex = 0;
        int maxCount = 0;
        Arrays.sort(initial);
        for(int i: initial) {
            // -满足仅感染一个，且是本集合代表元素
            if(count[i] > maxCount) {
                minIndex = i;
                maxCount = count[i];
            }
        }
        return minIndex;
    }

    // 创建简易并查集
    public static int MAX_N;
    public static int[] father;
    public static int[] size;
    public static int[] stack;
    // -添加额外标签：感染 infect 代表元素 representative
    // 感染 infect：true表示感染元素，false表示未感染元素
    // 代表元素 representative：-1表示集合中无感染元素，-2表示集合中存在两个以上感染元素，>=0表示一个感染元素
    public static boolean[] infect;
    public static int[] representative;

    public static void initialize(int n, int[] initial) {
        /**
         * 初始化并查集*/
        // 1 基本初始化
        MAX_N = n;
        father = new int[MAX_N];
        size = new int[MAX_N];
        stack = new int[MAX_N];
        infect = new boolean[MAX_N];
        representative = new int[MAX_N];
        for(int i = 0; i < MAX_N; i++) {
            father[i] = i;
            size[i] = 1;
            representative[i] = -1;
        }
        // 2 标签初始化
        for(int e: initial) {
            infect[e] = true;
            representative[e] = e;
        }
    }

    public static int find(int a) {
        /**
         * 找到a所在集合的代表元素+扁平化*/
        int i = 0;
        while(a != father[a]) {
            stack[i++] = a;
            a = father[a];
        }
        while(i > 0) {
            father[stack[--i]] = a;
        }
        return a;
    }

    public static boolean isSameSet(int a, int b) {
        /**
         * 判断a、b元素是否属于同一集合*/
        return find(a) == find(b);
    }

    public static void union(int a, int b) {
        /**
         * 合并a、b元素所在集合+小挂大*/
        int fa = find(a);
        int fb = find(b);
        if(fa != fb) {
            if(size[fa] >= size[fb]) {
                size[fa] += size[fb];
                father[fb] = fa;
            } else {
                size[fb] += size[fa];
                father[fa] = fb;
            }
        }
    }
}
